import torch
from transformers import BertModel, BertTokenizer

# 初始化BERT模型
bert_model = BertModel.from_pretrained('/gemini/data-2')
tokenizer = BertTokenizer.from_pretrained('/gemini/data-2')

# 构建输入
inputs = tokenizer.encode_plus('some input', return_tensors='pt') 

# 得到BERT模型输出 
outputs = bert_model(**inputs)

# pooled_output是序列的语义信息
pooled_output = outputs[1]  

# 第一个元素是[CLS] token对应的向量    
cls_output = outputs[0][:, 0, :]  

print(pooled_output.shape)   # [batch_size, hidden_size]
print(cls_output.shape)     # [batch_size, hidden_size]